package com.tqz.rc.server.push;

import com.tqz.rc.api.Cluster;
import com.tqz.rc.api.Instance;
import com.tqz.rc.api.ServiceInfo;
import com.tqz.rc.api.UdpPushAckPacket;
import com.tqz.rc.common.constant.Constants;
import com.tqz.rc.common.constant.DefaultValueConstant;
import com.tqz.rc.common.executor.NamedThreadFactory;
import com.tqz.rc.common.util.DateUtil;
import com.tqz.rc.common.util.JacksonUtils;
import com.tqz.rc.common.util.NamingUtils;
import com.tqz.rc.common.util.StringUtils;
import com.tqz.rc.server.core.Service;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.ApplicationListener;
import org.springframework.stereotype.Component;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetSocketAddress;
import java.net.SocketException;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.*;
import java.util.zip.GZIPOutputStream;

/**
 * <p>服务端通过udp方式向订阅的客户端推送服务变更的信息。
 * 比如`order-service`订阅了`product-service`，此时`product-service`下线了，服务端已经剔除了该服务，
 * 同时也要把该服务下线的通知发送给`order-service`，`order-service`再通过远程调用(dubbo/feign)`product-service`就会调用不到了。
 *
 * @author tianqingzhao
 * @since 2021/9/1 14:31
 */
@Component
public class UdpPushService implements ApplicationContextAware, ApplicationListener<ServiceChangeEvent> {

    private static final Logger log = LoggerFactory.getLogger(UdpPushService.class);

    private ApplicationContext applicationContext;

    private static DatagramSocket udpSocket;

    private static volatile ConcurrentMap<String, AckEntry> ackMap = new ConcurrentHashMap<>();

    private static volatile ConcurrentMap<String, Long> udpSendTimeMap = new ConcurrentHashMap<>();

    private static ConcurrentHashMap<String, Future> futureMap = new ConcurrentHashMap<>();

    private static ScheduledExecutorService executorService;

    private static ScheduledExecutorService udpSendExecutor;

    private DataSource dataSource = new DataSource() {

        @Override
        public String getData(Service service) {
            ServiceInfo result = new ServiceInfo();
            result.setClusters(getClusters(service));
            result.setCacheMillis(TimeUnit.SECONDS.toMillis(3));
            result.setLastRefTime(System.currentTimeMillis());
            result.setGroupName(service.getGroupName());
            result.setServiceName(NamingUtils.getGroupedName(service.getServiceName(), service.getGroupName()));
            result.setHosts(service.allIps(true));

            return JacksonUtils.toJson(result);
        }

        private String getClusters(Service service) {
            return StringUtils.join(service.getClusterMap().keySet(), ",");
        }
    };

    static {
        try {
            udpSocket = new DatagramSocket();

            Receiver receiver = new Receiver();
            Thread thread = new Thread(receiver);
            thread.setName("com.alibaba.nacos.naming.push.receiver");
            thread.setDaemon(true);
            thread.start();

            executorService = Executors.newScheduledThreadPool(1, new NamedThreadFactory("com.alibaba.nacos.naming.push.retransmitter"));

            udpSendExecutor = Executors.newScheduledThreadPool(1, new NamedThreadFactory("com.alibaba.nacos.naming.push.udpSender"));
        } catch (SocketException e) {
            e.printStackTrace();
        }
    }

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    @Override
    public void onApplicationEvent(ServiceChangeEvent event) {
        // 通过UDP方式将服务变动通知给订阅的客户端。
        Service service = event.getService();
        String namespaceId = service.getNamespaceId();
        String serviceName = service.getServiceName();

        String fullServiceName = assembleFullServiceName(namespaceId, serviceName);

        if (futureMap.containsKey(fullServiceName)) {
            return;
        }

        Future future = udpSendExecutor.schedule(() -> {
            try {
                // 获取该server下面所有的集群列表
                Map<String, Cluster> clusterMap = service.getClusterMap();
                for (Map.Entry<String, Cluster> entry : clusterMap.entrySet()) {
                    Cluster cluster = entry.getValue();

                    Set<Instance> ephemeralInstances = cluster.getEphemeralInstances();

                    AckEntry ackEntry = null;
                    long lastRefTime = System.nanoTime();
                    for (Instance instance : ephemeralInstances) {
                        // 向该服务的每个client端发送实例变更的udp推送，目前使用每个实例的端口号作为udp的端口号
                        ackEntry = prepareAckEntry(new InetSocketAddress(instance.getIp(), instance.getPort()), parseHostsData(service), lastRefTime);
                        System.out.println("向客户端：" + instance.getIp() + ",端口号：" + instance.getPort() + " 推送服务变更的信息");
                        udpPush(ackEntry);
                    }
                }
            } catch (Exception e) {

            } finally {
                futureMap.remove(fullServiceName);
            }
        }, 1000, TimeUnit.MILLISECONDS);

        futureMap.put(fullServiceName, future);
    }

    public static String assembleFullServiceName(String namespaceId, String serviceName) {
        return namespaceId + "##" + serviceName;
    }

    /**
     * 发送服务变更的事件
     *
     * @param service 变更的服务
     */
    public void serviceChanged(Service service) {
        this.applicationContext.publishEvent(new ServiceChangeEvent(this, service));
    }

    private static AckEntry prepareAckEntry(InetSocketAddress socketAddress, Map<String, Object> data,
                                            long lastRefTime) {
        if (data == null || data.size() <= 0) {
            log.error("[TqzRc-PUSH] pushing empty data for client is not allowed: {}", socketAddress);
            return null;
        }
        data.put("lastRefTime", lastRefTime);
        String dataStr = JacksonUtils.toJson(data);
        try {
            byte[] dataBytes = dataStr.getBytes(Constants.CHARSET_UTF_8);
            dataBytes = compressIfNecessary(dataBytes);
            return prepareAckEntry(socketAddress, dataBytes, data, lastRefTime);
        } catch (Exception e) {
            log.error("[TqzRc-PUSH] failed to compress data: {} to client: {}, error: {}", data, socketAddress, e);
            return null;
        }
    }

    private static AckEntry prepareAckEntry(InetSocketAddress socketAddress, byte[] dataBytes, Map<String, Object> data,
                                            long lastRefTime) {
        String key = AckEntry.getAckKey(socketAddress.getAddress().getHostAddress(), socketAddress.getPort(), lastRefTime);
        try {
            DatagramPacket packet = new DatagramPacket(dataBytes, dataBytes.length, socketAddress);
            AckEntry ackEntry = new AckEntry(key, packet);
            // we must store the key be fore send, otherwise there will be a chance the
            // ack returns before we put in
            ackEntry.setData(data);
            return ackEntry;
        } catch (Exception e) {
            log.error("[TqzRC-PUSH] failed to prepare data: {} to client: {}, error: {}", data, socketAddress, e);
        }
        return null;
    }

    /**
     * 解析实例信息
     *
     * @param service
     * @return
     */
    private Map<String, Object> parseHostsData(Service service) {
        String dataContent = dataSource.getData(service);
        Map<String, Object> result = new HashMap<String, Object>(2);
        result.put("type", "dom");
        result.put("data", dataContent);
        return result;
    }

    /**
     * 压缩数据
     *
     * @param dataBytes
     * @return
     * @throws IOException
     */
    private static byte[] compressIfNecessary(byte[] dataBytes) throws IOException {
        // 数据超过1kb进行压缩
        int maxDataSizeUncompress = 1024;
        if (dataBytes.length < maxDataSizeUncompress) {
            return dataBytes;
        }

        ByteArrayOutputStream out = new ByteArrayOutputStream();
        GZIPOutputStream gzip = new GZIPOutputStream(out);
        gzip.write(dataBytes);
        gzip.close();

        return out.toByteArray();
    }

    /**
     * 通过udp的方法发送推送数据
     *
     * @param ackEntry 数据
     */
    private void udpPush(AckEntry ackEntry) {
        if (ackEntry.getRetryTimes() > DefaultValueConstant.UDP_MAX_RETRY_TIMES) {
            log.warn("max re-push times reached, retry times {}, key: {}", ackEntry.getRetryTimes(),
                    ackEntry.getKey());
            ackMap.remove(ackEntry.getKey());
            udpSendTimeMap.remove(ackEntry.getKey());
            return;
        }

        try {
            ackMap.put(ackEntry.getKey(), ackEntry);
            udpSendTimeMap.put(ackEntry.getKey(), System.currentTimeMillis());

            udpSocket.send(ackEntry.getOrigin());
            ackEntry.increaseRetryTime();

            // 重试机制，实现的效果就是当前发送完udp之后，如果没有收到ACK响应，就隔10秒重发一次，并且只重试一次
//            executorService.schedule(new Retransmitter(ackEntry),
//                    TimeUnit.NANOSECONDS.toMillis(DefaultValueConstant.ACK_TIMEOUT_NANOS), TimeUnit.MILLISECONDS);
        } catch (IOException e) {
            e.printStackTrace();
            log.error("向客户端推送数据失败。");
        }
    }

    /**
     * 重试任务类
     */
    public class Retransmitter implements Runnable {

        AckEntry ackEntry;

        public Retransmitter(AckEntry ackEntry) {
            this.ackEntry = ackEntry;
        }

        @Override
        public void run() {
            if (ackMap.containsKey(ackEntry.getKey())) {
                log.info("retry to push data, retrytimes： " + ackEntry.getRetryTimes() + ",key: " + ackEntry.getKey());
                udpPush(ackEntry);
            }
        }
    }

    /**
     * 该内部类用于接收 `nacos客户端处理完服务端的推送` 的响应。
     * 也就是服务端通过udp推送给客户端，客户端处理完服务的变更之后，
     * 同时要再通过udp的方式给服务端一个响应。
     */
    public static class Receiver implements Runnable {

        @Override
        public void run() {
            log.info("服务端准备接收到客户端处理服务变更的响应，{}", DateUtil.getCurDateTime());
            while (true) {
                byte[] buffer = new byte[Constants.UDP_MSS];
                DatagramPacket packet = new DatagramPacket(buffer, buffer.length);

                try {
                    udpSocket.receive(packet);

                    String json = new java.lang.String(packet.getData(), 0, packet.getLength(), Constants.CHARSET_UTF_8).trim();
                    UdpPushAckPacket ackPacket = JacksonUtils.toObj(json, UdpPushAckPacket.class);
                    log.info("服务端已接收到客户端处理服务变更的响应，{}", DateUtil.getCurDateTime());

                    InetSocketAddress socketAddress = (InetSocketAddress) packet.getSocketAddress();
                    String ip = socketAddress.getAddress().getHostAddress();
                    int port = socketAddress.getPort();

                    String ackKey = AckEntry.getAckKey(ip, port, ackPacket.lastRefTime);
                    AckEntry ackEntry = ackMap.remove(ackKey);
                    if (ackEntry == null) {
                        throw new IllegalStateException(
                                "unable to find ackEntry for key: " + ackKey + ", ack json: " + json);
                    }

                    long pushCost = System.currentTimeMillis() - udpSendTimeMap.get(ackKey);

                    log.info("received ack: {} from: {}:{}, cost: {} ms, unacked: {}", json, ip, port, pushCost, ackMap.size());

                    udpSendTimeMap.remove(ackKey);
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }
}
